{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 二维声波问题\n",
    "\n",
    "[![下载Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_notebook.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/zh_cn/cfd_solver/mindspore_acoustic.ipynb)&emsp;[![下载样例代码](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_download_code.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/zh_cn/cfd_solver/mindspore_acoustic.py)&emsp;[![查看源文件](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_source.png)](https://gitee.com/mindspore/docs/blob/master/docs/mindflow/docs/source_zh_cn/cfd_solver/acoustic.ipynb)\n",
    "\n",
    "## 环境安装\n",
    "\n",
    "本案例要求 **MindSpore >= 2.4.0** 版本。具体请查看 [MindSpore安装](https://www.mindspore.cn/install)。\n",
    "\n",
    "此外，你需要安装 **MindFlow >=0.4.0** 版本。如果当前环境还没有安装，请按照下列方式选择后端和版本进行安装。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mindflow_version = \"0.4.0\"  # update if needed\n",
    "\n",
    "# Only NPU is supported.\n",
    "!pip uninstall -y mindflow-ascend\n",
    "!pip install mindflow-ascend==$mindflow_version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 概述\n",
    "\n",
    "声波方程求解是医疗超声、地质勘探等领域中的核心技术，大规模声波方程求解面临算力和存储的挑战。声波方程求解器一般采用频域求解算法和时域求解算法，时域求解算法的代表是时域有限差分法 (TDFD)，频域求解算法包括频域有限差分法 (FDFD)、有限元法 (FEM) 和 CBS (Convergent Born series) 迭代法。CBS 方法由于不引入频散误差，且求解的内存需求低，因此受到工程和学术界的广泛关注。尤其是 [Osnabrugge et al. (2016)](https://linkinghub.elsevier.com/retrieve/pii/S0021999116302595) 解决了该方法的收敛性问题，使得 CBS 方法的应用具有更广阔的前景。\n",
    "\n",
    "本案例将演示如何调用 MindFlow 提供的 CBS API 实现二维/三维声波方程的求解。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 问题描述\n",
    "\n",
    "声波方程求解中，波速场和震源信息是输入参数，求解输出的是时空分布的波场。\n",
    "\n",
    "二维/三维声波方程的表达式如下\n",
    "\n",
    "| 时域表达式                                            | 频域表达式                                        |\n",
    "| ----------------------------------------------------- | ------------------------------------------------- |\n",
    "| $\\frac{\\partial^2u}{\\partial t^2} - c^2 \\Delta u = f$ | $-\\omega^2 \\hat{u} - c^2 \\Delta\\hat{u} = \\hat{f}$ |\n",
    "\n",
    "其中\n",
    "\n",
    "- $u(\\bold{x},t) \\;\\; [L]$ 变形位移 (压强除以密度)，标量\n",
    "- $c(\\bold{x}) \\;\\; [L/T]$ 波速，标量\n",
    "- $f(\\bold{x},t) \\;\\; [L/T^2]$ 震源激励 (体积分布力)，标量\n",
    "\n",
    "实际求解中，为了降低参数维度，一般先将参数无量纲化，然后针对无量纲方程和参数进行求解，最后恢复解的量纲。选取 $\\omega$、$\\hat{f}$ 和 $d$（网格间距，本案例要求网格在各方向间距相等）对频域方程做无量纲化，可得频域无量纲方程：\n",
    "\n",
    "$$\n",
    "u^* + c^{*2} \\tilde{\\Delta} + f^* = 0\n",
    "$$\n",
    "\n",
    "其中\n",
    "\n",
    "- $u^* = \\hat{u} \\omega^2 / \\hat{f}$ 为无量纲变形位移\n",
    "- $c^* = c / (\\omega d)$ 为无量纲波速\n",
    "- $\\tilde{\\Delta}$ 为归一化 Laplace 算子，即网格间距均为 1 时的 Laplace 算子\n",
    "- $f^*$ 为标记震源位置的 mask，即在震源作用点为 1，其余位置为 0\n",
    "\n",
    "本案例中 `src` 包可以在 [src](https://gitee.com/mindspore/mindscience/tree/master/MindFlow/applications/cfd/acoustic/src) 下载。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import mindspore as ms\n",
    "from mindspore import Tensor\n",
    "\n",
    "from mindflow.utils import load_yaml_config\n",
    "\n",
    "from cbs.cbs import CBS\n",
    "from src import visual\n",
    "from solve_acoustic import solve_cbs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 定义输入参数和输出采样方式\n",
    "\n",
    "本案例所需的输入为有量纲 2D/3D 速度场、震源位置列表、震源波形，在文件 [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) (2D) or [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) (3D) 中指定输入文件名。为了方便用户直接验证，本案例在本[链接](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/cfd/acoustic)中提供了预置的输入数据，请下载所需要的数据集，并保存在 `./dataset` 目录下。2D数据集包括速度场 `velocity_2d.npy`、震源位置列表 `srclocs_2d.csv`、震源波形 `srcwaves_2d.csv`，3D数据集包括速度场 `velocity_3d.npy`、震源位置列表 `srclocs_3d.csv`、震源波形 `srcwaves_3d.csv`。用户可仿照输入文件格式自行修改输入参数。\n",
    "\n",
    "输出为时空分布的波场，为了明确输出如何在时间和频率上采样，需在 [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) 或 [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) 文件中指定 `dt`, `nt` 等参数。\n",
    "\n",
    "由于输入的震源波形在时间上的采样率可能与输出所要求的不一致，因此需对其进行插值。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ms.set_context(device_target='Ascend', device_id=0, mode=ms.GRAPH_MODE)\n",
    "\n",
    "# Space dimension\n",
    "dim = 2\n",
    "\n",
    "if dim == 2:\n",
    "    config = load_yaml_config(\"./config_2d.yaml\")\n",
    "else:\n",
    "    config = load_yaml_config(\"./config_3d.yaml\")\n",
    "\n",
    "data_config = config['data']\n",
    "solve_config = config['solve']\n",
    "summary_config = config['summary']\n",
    "\n",
    "# Coarsen rate for reducing scale\n",
    "rate = solve_config['coarsen_rate']\n",
    "\n",
    "# Read time & frequency points\n",
    "dt = solve_config['dt']\n",
    "nt = solve_config['nt']\n",
    "dt *= rate # Increase the time iteraion step size\n",
    "nt //= rate\n",
    "receivers = None\n",
    "\n",
    "# Read velocity array\n",
    "velo = np.load(os.path.join(data_config['root_dir'], data_config['velocity_field']))\n",
    "assert dim == len(velo.shape), \"dim and dimension of velocity should be equal\"\n",
    "dx = data_config['velocity_dx']\n",
    "dz = data_config['velocity_dz']\n",
    "if dim == 2:\n",
    "    dxs = (rate*dz, rate*dx)\n",
    "    velo = velo[::rate, ::rate] # Coarsen velocity field\n",
    "else:\n",
    "    dy = data_config['velocity_dy']\n",
    "    dxs = (rate*dz, rate*dy, rate*dx)\n",
    "    velo = velo[::rate, ::rate, ::rate] # Coarsen velocity field\n",
    "\n",
    "# Read source locations\n",
    "df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_locations']), index_col=0)\n",
    "if dim == 2:\n",
    "    slocs = df[['y', 'x']].values # Shape (ns, 2)\n",
    "else:\n",
    "    slocs = df[['z', 'y', 'x']].values # Shape (ns, 3)\n",
    "\n",
    "# Read & interp source wave\n",
    "df = pd.read_csv(os.path.join(data_config['root_dir'], data_config['source_wave']))\n",
    "inter_func = interp1d(df.t, df.f, kind='cubic', bounds_error=False, fill_value=0)\n",
    "\n",
    "ts = np.arange(nt) * dt\n",
    "omegas_all = np.fft.rfftfreq(nt) * (2 * np.pi / dt)\n",
    "\n",
    "# Interpolation source wave\n",
    "src_waves = inter_func(ts) # Shape (nt)\n",
    "src_amplitudes = np.fft.rfft(src_waves) # Shape (nt//2+1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 选取待求频点\n",
    "\n",
    "确定了输出采样方式即确定了所有待求频点。但为了减少计算量，也可以只选择部分频点进行求解，其余频点通过插值获得。具体的频点降采样方式由 [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) 和 [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) 文件中的 `downsample_mode`, `downsample_rate` 指定。默认为不做降采样，即求解除 $\\omega=0$ 之外的所有频点。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select omegas\n",
    "no = len(omegas_all) // solve_config['downsample_rate']\n",
    "\n",
    "if solve_config['downsample_mode'] == 'exp':\n",
    "    omegas_sel = np.exp(np.linspace(np.log(omegas_all[1]), np.log(omegas_all[-1]), no))\n",
    "elif solve_config['downsample_mode'] == 'square':\n",
    "    omegas_sel = np.linspace(omegas_all[1]**.5, omegas_all[-1]**.5, no)**2\n",
    "else:\n",
    "    omegas_sel = np.linspace(omegas_all[1], omegas_all[-1], no)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 执行仿真\n",
    "\n",
    "将相关数组定义为 Tensor，调用 `solve_cbs()`，在 NPU 执行求解。由于显存限制，求解过程在频点维度分批执行，batch 数量由用户在 `config_2d.yaml` 或 `config_3d.yaml` 中指定，不要求整除频点数（允许最后一个 batch 的 size 与其他 batch 不一致）。各方向上的吸收边界厚度和边界类型可分别通过 [config_2d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_2d.yaml) 或 [config_3d.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/cfd/acoustic/config_3d.yaml) 文件中的 `pml_size`, `btype` 来指定。求解完成后，会保存频域求解结果到文件 `u_star.npy`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Send to NPU and perform computation\n",
    "os.makedirs(summary_config['root_dir'], exist_ok=True)\n",
    "velo = Tensor(velo, dtype=ms.float32, const_arg=True)\n",
    "\n",
    "dxs_nd = tuple(d / dxs[-1] for d in dxs) # Nondimensional dxs\n",
    "pml_size = tuple(solve_config['pml_size'])\n",
    "btype = solve_config['btype']\n",
    "cbs = CBS(velo.shape, dxs=dxs_nd, pml_size=pml_size, alpha=1., rampup=12,\n",
    "          btype=btype, remove_pml=False)\n",
    "\n",
    "# Shape (ns, no, len(receiver_zs), nx) or (ns, no, len(receiver_zs), ny, nx)\n",
    "ur, ui, errs = solve_cbs(\n",
    "    cbs, velo, slocs, omegas_sel, receivers=receivers, dxs=dxs, n_batches=solve_config['n_batches'])\n",
    "\n",
    "u_star = ur.numpy() + 1j * ui.numpy() # Shape (ns, no, len(krs), nx) or (ns, no, len(krs), ny, nx)\n",
    "\n",
    "np.save(os.path.join(summary_config['root_dir'], 'u_star.npy'), np.squeeze(u_star))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 仿真结果后处理\n",
    "\n",
    "CBS 求解的是无量纲化的频域方程，但下游任务通常希望在时域观察有量纲波场的演化过程，因此最后将求解结果恢复量纲化并转回时域。恢复量纲的方式为 $\\hat{u} = u^* \\hat{f} / \\omega^2$，若在前述的“选取待求频点”步骤中对频点做了降采样，则在此处需沿频率方向插值恢复所有频点的解。然后对有量纲的频域波场 $\\hat{u}$ 做 Fourier 反变换得到时域波场 $u$。将时域波场保存至文件 `u_time.npy`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Recover dimension and interpolate to full frequency domain\n",
    "ones = (1,) * dim\n",
    "u_star /= omegas_sel.reshape(-1, *ones)**2\n",
    "u_star = interp1d(omegas_sel, u_star, axis=1, kind='cubic', bounds_error=False, fill_value=0)(omegas_all)\n",
    "u_star *= src_amplitudes.reshape(-1, *ones)\n",
    "\n",
    "# Transform to time domain\n",
    "u_time = np.fft.irfft(u_star, axis=1)\n",
    "np.save(os.path.join(summary_config['root_dir'], 'u_time.npy'), u_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，读取时域波场并可视化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the result\n",
    "u_time = np.load(os.path.join(summary_config['root_dir'], 'u_time.npy'))\n",
    "if dim == 2:\n",
    "    visual.anim(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))\n",
    "else:\n",
    "    visual.anim3d(velo.numpy(), u_time, ts, os.path.join(summary_config['root_dir'], 'wave.gif'))\n",
    "visual.plot_errs(errs, os.path.join(summary_config['root_dir'], 'errors.png'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6 (main, Nov 14 2022, 16:10:14) [GCC 11.3.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
